/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef OPS_BUILD_IN_OP_TILING_RUNTIME_UNSORTED_SEGMENT_SUM_H
#define OPS_BUILD_IN_OP_TILING_RUNTIME_UNSORTED_SEGMENT_SUM_H

#include "register/op_compile_info_base.h"
#include "auto_tiling_rt2.h"
#include "vector_tiling_rt2.h"

namespace optilingUss {
struct TilingPrepareForUnsortedSegmentSumCompileInfo {
    std::shared_ptr<TbeReduce::AutoTilingCompileInfo> dsl_compile_info;
    int32_t core_num;
    int32_t ub_size;
    int32_t ub_tensor_num;
    int32_t impl_mode;
    bool is_tik{ false };
};

template <typename T> bool ReadCompileItem(const nlohmann::json &all_vars, const std::string &name, T &value)
{
    if (all_vars.empty()) {
        return false;
    }

    if (all_vars.count(name) == 0) {
        return false;
    }

    value = all_vars[name].get<T>();
    return true;
}

template <typename T1, typename T2>
void ReadCompileItem(const nlohmann::json &all_vars, const std::string &name, T1 &value, const T2 default_value)
{
    if (!ReadCompileItem(all_vars, name, value)) {
        value = static_cast<T1>(default_value);
    }
}

struct UnsortedSegmentSumTilingParamsFp32 {
    // common params
    int32_t select_key_input_scalar;
    int32_t need_core_num_input_scalar;

    // input data params
    // front core
    int32_t input_ele_num_front_core_input_scalar;
    // front part front core
    int32_t input_mov_times_gm2ub_front_part_front_core_input_scalar;
    int32_t input_front_burst_len_front_part_front_core_input_scalar;
    int32_t input_last_burst_len_front_part_front_core_input_scalar;
    int32_t input_front_ele_num_ub_front_part_front_core_input_scalar;
    int32_t input_last_ele_num_ub_front_part_front_core_input_scalar;
    int32_t input_front_rows_front_part_front_core_input_scalar;
    int32_t input_last_rows_front_part_front_core_input_scalar;
    // last part front core
    int32_t input_mov_times_gm2ub_last_part_front_core_input_scalar;
    int32_t input_front_burst_len_last_part_front_core_input_scalar;
    int32_t input_last_burst_len_last_part_front_core_input_scalar;
    int32_t input_front_ele_num_ub_last_part_front_core_input_scalar;
    int32_t input_last_ele_num_ub_last_part_front_core_input_scalar;
    int32_t input_front_rows_last_part_front_core_input_scalar;
    int32_t input_last_rows_last_part_front_core_input_scalar;
    // last core
    int32_t input_ele_num_last_core_input_scalar;
    // front part last core
    int32_t input_mov_times_gm2ub_front_part_last_core_input_scalar;
    int32_t input_front_burst_len_front_part_last_core_input_scalar;
    int32_t input_last_burst_len_front_part_last_core_input_scalar;
    int32_t input_front_ele_num_ub_front_part_last_core_input_scalar;
    int32_t input_last_ele_num_ub_front_part_last_core_input_scalar;
    int32_t input_front_rows_front_part_last_core_input_scalar;
    int32_t input_last_rows_front_part_last_core_input_scalar;
    // last part last core
    int32_t input_mov_times_gm2ub_last_part_last_core_input_scalar;
    int32_t input_front_burst_len_last_part_last_core_input_scalar;
    int32_t input_last_burst_len_last_part_last_core_input_scalar;
    int32_t input_front_ele_num_ub_last_part_last_core_input_scalar;
    int32_t input_last_ele_num_ub_last_part_last_core_input_scalar;
    int32_t input_front_rows_last_part_last_core_input_scalar;
    int32_t input_last_rows_last_part_last_core_input_scalar;

    // e num params
    int32_t e_num_input_scalar;
    int32_t e_mov_times_gm2ub_input_scalar;
    int32_t e_ub2gm_front_burst_len_input_scalar;
    int32_t e_num_front_part_input_scalar;
    int32_t e_ub2gm_last_burst_len_input_scalar;
    int32_t e_num_last_part_input_scalar;

    // ids params
    int32_t ids_size_input_scalar;
    int32_t ids_ele_num_front_core_input_scalar;
    int32_t ids_mov_times_gm2ub_front_core_input_scalar;
    int32_t ids_front_burst_len_front_core_input_scalar;
    int32_t ids_last_burst_len_front_core_input_scalar;
    int32_t ids_ele_num_ub_front_part_front_core_input_scalar;
    int32_t ids_ele_num_ub_last_part_front_core_input_scalar;
    int32_t ids_ele_num_last_core_input_scalar;
    int32_t ids_mov_times_gm2ub_last_core_input_scalar;
    int32_t ids_front_burst_len_last_core_input_scalar;
    int32_t ids_last_burst_len_last_core_input_scalar;
    int32_t ids_ele_num_ub_front_part_last_core_input_scalar;
    int32_t ids_ele_num_ub_last_part_last_core_input_scalar;

    // output init params
    int32_t output_ub_init_last_repeat_time_front_part_front_core_input_scalar;
    int32_t output_ub_init_times_front_part_front_core_input_scalar;
    int32_t output_ub_init_last_repeat_time_last_part_front_core_input_scalar;
    int32_t output_ub_init_times_last_part_front_core_input_scalar;
    int32_t output_ub_init_last_repeat_time_front_part_last_core_input_scalar;
    int32_t output_ub_init_times_front_part_last_core_input_scalar;
    int32_t output_ub_init_last_repeat_time_last_part_last_core_input_scalar;
    int32_t output_ub_init_times_last_part_last_core_input_scalar;
    int32_t input_last_axis_align_front_part_ele_num_input_scalar;
    int32_t input_last_axis_align_floor_ele_num_input_scalar;
    int32_t last_part_vadd_mask_input_scalar;
    int32_t e_gm2ub_last_burst_len_input_scalar;
    int32_t output_ub_init_last_row_last_repeat_time_front_part_front_core_input_scalar;
    int32_t output_ub_init_last_row_times_front_part_front_core_input_scalar;
    int32_t output_ub_init_last_row_last_repeat_time_last_part_front_core_input_scalar;
    int32_t output_ub_init_last_row_times_last_part_front_core_input_scalar;
    int32_t output_ub_init_last_row_last_repeat_time_front_part_last_core_input_scalar;
    int32_t output_ub_init_last_row_times_front_part_last_core_input_scalar;
    int32_t output_ub_init_last_row_last_repeat_time_last_part_last_core_input_scalar;
    int32_t output_ub_init_last_row_times_last_part_last_core_input_scalar;

    int32_t num_segments;

    int32_t tiling_core_num;
    int32_t repeat_front_front_part_front_core;
    int32_t col_sub_block_front_front_part_front_core;
    int32_t repeat_last_front_part_front_core;
    int32_t col_sub_block_last_front_part_front_core;
    int32_t repeat_front_last_part_front_core;
    int32_t col_sub_block_front_last_part_front_core;
    int32_t repeat_last_last_part_front_core;
    int32_t col_sub_block_last_last_part_front_core;
    int32_t repeat_front_front_part_last_core;
    int32_t col_sub_block_front_front_part_last_core;
    int32_t repeat_last_front_part_last_core;
    int32_t col_sub_block_last_front_part_last_core;
    int32_t repeat_front_last_part_last_core;
    int32_t col_sub_block_front_last_part_last_core;
    int32_t repeat_last_last_part_last_core;
    int32_t col_sub_block_last_last_part_last_core;
    int32_t e_num_sub;
    int32_t vadd_repeat_255;
    int32_t vadd_repeat_64;
    int32_t vadd_repeat_last;
    int32_t move_pad;
    int32_t max_cache_n_num;
    int32_t repeat_remove_pad;
    int32_t col_block_remove_pad;
    int32_t cache_num_block;
};

struct UnsortedSegmentSumTilingParamsInt32 {
    // common params
    int32_t select_key_input_scalar;
    int32_t need_core_num_input_scalar;
    int32_t num_segments_front_core_input_scalar;
    int32_t num_segments_last_core_input_scalar;

    // ids params
    int32_t ids_size_input_scalar;
    int32_t ids_mov_times_gm2ub_input_scalar;
    int32_t ids_ele_num_ub_front_part_input_scalar;
    int32_t ids_front_burst_len_input_scalar;
    int32_t ids_ele_num_ub_last_part_input_scalar;
    int32_t ids_last_burst_len_input_scalar;

    // e num params
    int32_t e_num_input_scalar;
    int32_t e_mov_times_gm2ub_input_scalar;
    int32_t e_ub2gm_front_burst_len_input_scalar;
    int32_t e_num_front_part_input_scalar;
    int32_t repeat_time_front_part_input_scalar;
    int32_t e_ub2gm_last_burst_len_input_scalar;
    int32_t e_num_last_part_input_scalar;
    int32_t repeat_time_last_part_input_scalar;
    int32_t align_scalar;
    int32_t align_scalar_lastcore;

    int32_t e_gm2ub_front_burst_len_input_scalar;
    int32_t e_gm2ub_last_burst_len_input_scalar;
    int32_t num_segment_max;
    int32_t num_segment_max_time;
    int32_t num_segment_max_time_lastcore;
    int32_t front_num_segment;
    int32_t front_num_segment_last;
    int32_t front_num_segment_lastcore;
    int32_t front_num_segment_last_lastcore;
    int32_t e_ub2gm_front_burst_len_input_scalar_lastcore;
    int32_t e_ub2gm_last_burst_len_input_scalar_lastcore;
    int32_t repeat_times;
    int32_t repeat_times_last_part;
    int32_t repeat_times_last_part_lastcore;
    int32_t e_mov_times_gm2ub_input_scalar_lastcore;
    int32_t repeat_time_front_part_input_scalar_lastcore;
    int32_t num_segments;

    int32_t tiling_core_num;
};

ge::graphStatus TilingPrepare4SegmentSumComm(gert::TilingParseContext *context);
ge::graphStatus Tiling4UnsortedSegmentSum(gert::TilingContext *context);
} // namespace optilingUss
#endif // OPS_BUILD_IN_OP_TILING_RUNTIME_UNSORTED_SEGMENT_SUM_H
